#!/usr/bin/env python3
"""
Tool which is used for merge all the files that included depth JPEGImages mask gt.yml info.yml and intrinsics.json 
Author: cui
"""
from hashlib import new
from itertools import count
import os
import cv2
import tqdm
import json
import yaml
import numpy as np
import glob
import sys
import json



########################---main---########################
def make_directories(folder):
    if not os.path.exists(folder+"JPEGImages/"):
        os.makedirs(folder+"JPEGImages/")
    if not os.path.exists(folder+"depth/"):
        os.makedirs(folder+"depth/")
    if not os.path.exists(folder+"mask/"):
        os.makedirs(folder+"mask/")

def read_and_write_jpg(input_folder,output_folder,input_idx,output_index):
    path_JPEGImages = input_folder+"{}.jpg".format(input_idx)
    img_rgb = cv2.imread(path_JPEGImages)
    # print("input_idx:",input_idx,"output_idx: ",output_index)
    cv2.imwrite(output_folder+'{}.jpg'.format(output_index),img_rgb)

def read_and_write_png(input_folder,output_folder,input_idx,output_index):
    path_JPEGImages = input_folder+"{}.png".format(input_idx)
    img_rgb = cv2.imread(path_JPEGImages)
    # print("input_idx:",input_idx,"output_idx: ",output_index)
    cv2.imwrite(output_folder+'{}.png'.format(output_index),img_rgb)

def read_yaml(input_folder):
    with open(input_folder, 'r') as scene:
        temp_data = yaml.load(scene,Loader=yaml.FullLoader)
        return temp_data

def print_usage():
    print("Usage: merge_files.py <path>")
    print("path: all or name of the folder")
    print("e.g., python merge_files.py all merge")
    print("merge is what you want to generate.")
    
    
if __name__ == "__main__":
    try:
        if sys.argv[1] == "all":
            folders = glob.glob("LINEMOD/0*/")
            # print(folders)
            folder_out = "LINEMOD/"+sys.argv[2]+"/"
            # print(folder_out)
        # elif sys.argv[2:-1] in glob.glob("LINEMOD/*/"):
        #     a = sys.argv[2:-1]
        #     print(a)
        #     folders = [sys.argv[2:-1]+"/"]
        #     folder_out = sys.argv[-1]+"/"
        else:
            # print_usage()
            exit()
    except:
        print_usage()
        exit()

#############  output files   ############
    make_directories(folder_out)
    path_JPEGImages_out = folder_out+"JPEGImages/"
    path_depth_out = folder_out+"depth/"
    path_mask_out = folder_out+"mask/"
    path_gtyaml_out = folder_out+"gt.yml"
    path_infoyaml_out = folder_out+"info.yml"
    # print(path_JPEGImages_out,path_depth_out,path_mask_out)

#############  Total floders num   ############
    folders_num = len(folders)
    # print("Total floders num is :",folders_num)
    # print("folders_list: ",folders)
     
    # total_img_num = 0
    # for k ,l in enumerate(folders):
    #     path_JPEGImages_folder = os.path.join(l,"JPEGImages/")
    #     files_num = len(os.listdir(path_JPEGImages_folder)) 
    #     total_img_num = total_img_num+files_num
    # print("total_img_num: ",total_img_num)


#############  merge files   ############
    scene = {}
    scene_info = {}
    idx = 0
    for num,folder in enumerate(folders):
        print("num:",num,"folders:",folder)
        #############   input files path   ############
        path_JPEGImages_folder = os.path.join(folder,"JPEGImages/")
        path_mask_folder = os.path.join( folder,"mask/")
        path_depth_folder =   os.path.join(folder,"depth/")
        path_gtyaml =  os.path.join(folder,"gt.yml")
        path_infoyaml = os.path.join(folder,"info.yml")
        # print(path_JPEGImages_folder,path_mask_folder,path_depth_folder)
        
        ##############   move intrinsics.json   #################
        if num == 0:
                with open(folder+'intrinsics.json', 'r') as f:
                    camera_intrinsics = json.load(f)      
                with open(folder_out+'intrinsics.json', 'w') as fp:
                    json.dump(camera_intrinsics, fp)       

        files_num = len(os.listdir(path_JPEGImages_folder)) #Count the number of samples in this folder
        # print("files_num",files_num)
        temp_gtyaml_data = read_yaml(path_gtyaml)
        temp_infoyaml_data = read_yaml(path_infoyaml)
        for i in range(files_num):
            read_and_write_jpg(path_JPEGImages_folder,path_JPEGImages_out,i,idx)
            read_and_write_png(path_mask_folder,path_mask_out,i,idx)
            read_and_write_png(path_depth_folder,path_depth_out,i,idx)
            scene[idx] = temp_gtyaml_data[i]
            depth_scale_scene_info = temp_infoyaml_data[i]
            depth_scale_scene_info["depth_scale"] = 1
            scene_info[idx] = depth_scale_scene_info
        
            idx = idx+1
        print(len(scene))

    with open(path_gtyaml_out, 'w') as f:
        yaml.dump(scene, f)
    with open(path_infoyaml_out, 'w') as f:
        yaml.dump(scene_info, f)


